'''
mymodel.py

Copyright 2009 Xavier Mendez Navarro aka Javi

This file is part of pysqlin

pysqlin is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation version 2 of the License.

pysqlin is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with pysqlin; if not, write to the Free Software
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
'''

from framework.queries import Database
from framework.myexceptions import MyException
from framework.settings import Settings

from collections import defaultdict
from UserDict import UserDict
import urlparse
import urllib
import re
import pickle


class DataModel:
    def __init__(self, id):
	self.id = id
	self.filters_data = FilterData()
	self.kbase = KnowledgeBase()
	self.inj_data = InjectionData()

    def reset_kbase(self):
	self.kbase = KnowledgeBase()

    def save(self, filename):
	output = None
	try:
	    output = open(filename, 'w+b')
	    pickle.dump(self, output)
	except IOError, e:
	    raise MyException(e)
	finally:
	    if output: output.close()

    def load(self, path):
	pkl_file = None
	pkl = None

	try:
	    pkl_file = open(path, 'r+b')
	    pkl = pickle.load(pkl_file)
	except IOError, e:
	    raise MyException(e)
	except EOFError, e:
	    raise MyException(e)
	finally:
	    if pkl_file: pkl_file.close()

	return pkl

class FilterData(UserDict):
    def __init__(self, data = {}):
	self.data = data

    def update(self, plugin_data):
	for k, v in plugin_data.items():
	    self.data[k] = (v.enabled, v.priority)

    def merge(self, plugin_data):
	for k, v in [(k, v) for k, v in plugin_data.items() if not self.data.has_key(k)]:
	    self.data[k] = (v.enabled, v.priority)

class InjectionData(UserDict):
    ERROR_MODE_POSITIVE, ERROR_MODE_NEGATIVE = 'positive', 'negative'

    PLUGIN = 'use'
    URL = 'url'
    POST_DATA = 'post_data'
    COOKIE = 'cookie'
    DATABASE = 'database'
    AUTH_METHOD = 'http_auth_method'
    AUTH_CREDENTIALS = 'http_auth_cred'

    INJECTOR = 'injector'
    ERROR_MODE = 'error_mode'
    ERROR_STRING = 'error_string'

    def __init__(self):
	self.data = {}
	self.data['main'] = {
	    InjectionData.URL: ("", 'Injection URL'),
	    InjectionData.POST_DATA: ("", 'HTTP Post data'),
	    InjectionData.COOKIE: ("", 'HTTP Cookie value'),
	    InjectionData.ERROR_STRING: ("", 'Injection error string'),
	    InjectionData.ERROR_MODE: ("", 'Injection error string mode'),
	    InjectionData.AUTH_METHOD: ("", 'HTTP authentication mode'),
	    InjectionData.AUTH_CREDENTIALS: ("", 'HTTP authentication credentials'),
	    InjectionData.DATABASE: ("", 'Injection database'),
	    InjectionData.INJECTOR: ("", 'Injection exploitation plugin')
	}
	self.data['plugins'] = {}
	self.data['injectors'] = {}

    # ------------------------------------------------
    # Factory functions
    # ------------------------------------------------
    @staticmethod
    def from_raw(url, post_data, cookie, error_string, error_mode, auth_method, auth_cred, database):
	idata = InjectionData()
	idata.data = {}
	idata.data['plugins'] = {}
	idata.data['injectors'] = {}
	idata.data['main'] = {
	    InjectionData.URL: (url, 'Injection URL'),
	    InjectionData.POST_DATA: (post_data, 'HTTP Post data'),
	    InjectionData.COOKIE: (cookie, 'HTTP Cookie value'),
	    InjectionData.ERROR_STRING: (error_string, 'Injection error string'),
	    InjectionData.ERROR_MODE: (error_mode, 'Injection error string mode'),
	    InjectionData.AUTH_METHOD: (auth_method, 'HTTP authentication mode'),
	    InjectionData.AUTH_CREDENTIALS: (auth_cred, 'HTTP authentication credentials'),
	    InjectionData.DATABASE: (database, 'Injection database'),
	    InjectionData.INJECTOR: ("", 'Injection exploitation plugin'),
	}

	return idata

    # ------------------------------------------------
    # Merge data functions
    # ------------------------------------------------
    def merge_session(self, old):
	self.delete('plugins')
	self.merge('plugins', old.data['plugins'])

    def merge(self, subdata, new_data):
	if subdata not in self.data.keys():
	    raise KeyError

	for k in new_data.keys():
	    self.data[subdata][k] = new_data[k]

    def update(self, new_data):
	for subdata in new_data.data.keys():
	    for k in new_data.data[subdata].keys():
		self.data[subdata][k] = new_data.data[subdata][k]

    def update_from_rawsqli(self, db, var, method, itype):
	if db is None:
	    raise MyException('Cannot convert from raw sqli. Database not set.')

	self[InjectionData.DATABASE] = db

	if method == 'GET':
	    (scheme, netloc, path, query, fragment) = urlparse.urlsplit(self[InjectionData.URL])
	    #url_params = urlparse.parse_qs(query)
	    url_params = dict(part.split('=') for part in query.split('&'))

	    url_params[var] = url_params[var] + Database.tests[db][itype]

	    self[InjectionData.URL] = urllib.unquote(urlparse.urlunsplit((scheme, netloc, path, urllib.urlencode(url_params), fragment)))
	elif method == 'POST':
	    url_params = dict(part.split('=') for part in self[InjectionData.POST_DATA].split('&'))
	    url_params[var] = url_params[var] + Database.tests[db][itype]
	    self[InjectionData.POST_DATA] =  urllib.unquote(urllib.urlencode(url_params))

    def merge_plugin(self, plugin):
	self.delete('plugins')
	self.merge('plugins', plugin.parameters())

    # ------------------------------------------------
    # dictionary functions
    # ------------------------------------------------
    def __getitem__(self, key):
	for subdata in self.data.values():
	    if subdata.has_key(key):
		value, desc = subdata[key]
		return value

	raise KeyError

    def __setitem__(self, key, item):
	for subdata in self.data.values():
	    if subdata.has_key(key):
		value, desc = subdata[key]
		subdata[key] = (item, desc)
		return

	raise KeyError

    def __delitem__(self, key): 
	for subdata in self.data.values():
	    if subdata.has_key(key):
		del subdata[key]
		return

	raise KeyError

    def delete(self, subdata):
	if subdata not in self.data.keys():
	    raise KeyError

	self.data[subdata] = {}

    def has_key(self, key):
	for subdata in self.data.values():
	    if subdata.has_key(key):
		return True
	return False

    def values(self):
	l = []
	for subdata in self.data.values():
	    l = l + subdata.values()
	return l

    def keys(self):
	l = []
	for subdata in self.data.values():
	    l = l + subdata.keys()
	return l

    def items(self):
	l = []
	l2 = []
	for subdata in self.data.values():
	    l = l + subdata.keys()
	    l2 = l2 + [value for value, desc in subdata.values()]
	return zip(l,l2)

    def items2(self):
	l = {}
	for key, subdata in self.data.items():
	    l[key] = [[name, item[0], item[1]] for name, item in subdata.items()]
	return l


class KnowledgeBase:
    def __init__(self):
	self.__data = defaultdict(list)

    def get(self, item):
	return self.__data[item]

    def keys(self):
	return self.__data.keys()

    def add(self, query, result):
	if not result:
	    return

	values = re.split(query.split_string, result)
	fields = query.expand_kbase(values)

	if len(fields) > 0:
	    if len(values) != len(fields):
		raise MyException('Knowledge fields mismatch, check injections.xml file.')

	    where = zip(values, fields)

	    for value, field in where:
		if field: 
		    if self.__data.has_key(field):
			if value not in self.__data[field]: self.__data[field].append(value)
		    else:
			self.__data[field].append(value)
